from django.shortcuts import render
from flybot.settings import PUBLIC_KEY, BASE_DIR
from django.http import JsonResponse
from django.db.models import Q
from flybot.settings import MJ_DEFAULT_CONFIG, COS_CONFIG
from flybotServer.utils.deploy_utils import check_env, deploy_wx_gpt, deploy_wx_mj
from flybotServer.utils.translate_utils import BaiduTranslate
from flybotServer.utils.oss_utils import Cos
from flybotServer.utils.midjourney_utils import Sender, Receiver
from flybotServer.models import BotApiKey, BotInfo, BotMessage, BotActions, MJTasks, SensitiveWord
from threading import Thread
from PIL import Image, ImageOps
import json, time, re, requests, os, openai

# 定义全局变量
# mj 机器人接收器开关
MJ_RECEIVER_ON = False
# 回应生成器开关
GENERATE_ACTIONS_ON = False
# 敏感词初始化
sensitiveWords = SensitiveWord.objects.all()
# 百度翻译初始化
baidu = BaiduTranslate()
# 对象存储初始化
cos = Cos(secret_id=COS_CONFIG['secret_id'], secret_key=COS_CONFIG['secret_key'], region=COS_CONFIG['region'])

# 机器人远程部署接口
def deploy(request):
    # 获取 api_key，如果过期返回过期
    api_key = request.headers.get("X-Api-Key")
    init_api_key_list = BotApiKey.objects.filter(state="init")
    is_api_key_used = True
    for botApiKey in init_api_key_list:
        if botApiKey.api_key == api_key:
            is_api_key_used = False

    if is_api_key_used:
        return JsonResponse({
            "code": 500,
            "message": "API_KEY 已过期"
        })

    api_url = json.loads(request.body).get("api_url")
    bot_type = json.loads(request.body).get("bot_type")
    hostname = json.loads(request.body).get("hostname")
    username = json.loads(request.body).get("username")
    password = json.loads(request.body).get("password")
    extra = json.loads(request.body).get("extra")

    # 检查环境待完善
    # if check_env(hostname, username, password):
    #     print("tester")

    QR_code = ""

    if bot_type == "wx_gpt":
        config = '''
        {
          "api_url": "http://firstfly.natapp4.cc",
          "api_key": "%s",
          "bot_type": "wx_gpt",
          "hostname": "%s",
          "username": "%s",
          "password": "%s",
          "extra": %s
        }
        '''%(api_key, hostname, username, password, json.dumps(extra))
        QR_code = deploy_wx_gpt(hostname, username, password, config)

    if bot_type == "wx_mj":
        config = '''
        {
          "api_url": "%s",
          "api_key": "%s",
          "bot_type": "wx_mj",
          "hostname": "%s",
          "username": "%s",
          "password": "%s",
          "extra": %s
        }
        ''' % (api_url, api_key, hostname, username, password, json.dumps(extra))
        QR_code = deploy_wx_mj(hostname, username, password, config)

    if bot_type == "qywx_gpt":
        pass

    if bot_type == "qywx_mj":
        pass

    return JsonResponse({
        "QR_code": QR_code
    })

# 机器人注册接口
def register(request):
    # 获取密钥信息
    api_key = request.headers.get("X-Api-Key")
    # 查询密钥信息
    try:
        botApiKey = BotApiKey.objects.get(api_key=api_key, state="init")
    except BotApiKey.DoesNotExist:
        # 如果密钥不存在，则返回错误信息
        return JsonResponse({
            "code": 400,
            "message": "无相关密钥或是密钥已经被激活"
        })
    else:
        # 如果密钥存在，则注册机器人相关信息
        botInfo = BotInfo()
        params = json.loads(request.body)
        botInfo.api_key = botApiKey
        botInfo.bot_type = params.get("bot_type", "")
        botInfo.hostname = params.get("hostname", "")
        botInfo.username = params.get("username", "")
        botInfo.password = params.get("password", "")
        botInfo.bot_path = params.get("bot_path", "")
        botInfo.process_id = params.get("process_id", "")
        botInfo.extra = params.get("extra", "")
        botInfo.save()
        botApiKey.state = "active"
        botApiKey.save()

    return JsonResponse({
        "code": 0
    })

# 机器人接收消息的回调接口
def message_subscriber(request):
    # 获取密钥信息
    api_key = request.headers.get("X-Api-Key")
    # 查询密钥信息
    try:
        botApiKey = BotApiKey.objects.get(api_key=api_key)
        botInfo = BotInfo.objects.get(api_key=botApiKey)
    except BotApiKey.DoesNotExist:
        # 如果密钥不存在，则返回错误信息
        return JsonResponse({
            "code": 400,
            "message": "无相关密钥或是密钥已经被激活"
        })
    except BotInfo.DoesNotExist:
        # 如果机器人信息不存在，则返回错误信息
        return JsonResponse({
            "code": 400,
            "message": "机器人未注册，请联系管理员"
        })
    else:
        botMessage = BotMessage()
        params = json.loads(request.body)
        botMessage.bot_info = botInfo
        botMessage.message = json.dumps(params.get("message", "{}"))
        botMessage.source = params.get("source", "{}")
        botMessage.state = "waiting"
        botMessage.save()
    return JsonResponse({
        "code": 0
    })

# 处理 mj 已发送任务并生成机器人消息
def process_sent_tasks():
    global MJ_RECEIVER_ON
    while MJ_RECEIVER_ON:
        # 获取已经发送的 mj 任务
        mjtasks = MJTasks.objects.filter(status="sent")[:10]
        for mjtask in mjtasks:
            # todo: 任务超过 X 分钟，将任务状态改为 start 并处理下一个
            if False:
                # 将状态改为已开始
                mjtask.status = "start"
                mjtask.save()
                continue
            extra = json.loads(mjtask.bot_info.extra.replace("'", "\""))
            # 创建一个接收器
            mjreceiver = Receiver(
                authorization=extra.get("authorization", MJ_DEFAULT_CONFIG["authorization"]),
                channelid=extra.get("channelid", MJ_DEFAULT_CONFIG["channelid"]),
                proxy="127.0.0.1:7890"
            )
            try:
                time.sleep(3)
                mjreceiver.collecting_results()
            except:
                # 有异常则跳过
                continue
            else:
                # 等待中的任务逻辑
                for awaiting_list_id in mjreceiver.awaiting_list.index:
                    prompt = mjreceiver.awaiting_list.loc[awaiting_list_id][0].replace("'", " ").strip()[0:1500]
                    print(prompt)
                    if mjtask.trans_prompt == prompt:
                        # 生成一条机器人消息，并将任务状态置为已开始
                        botActions = BotActions()
                        botActions.bot_info = mjtask.bot_info
                        botActions.action_name = "send_text"
                        botActions.action_content = "@{} \n🚀任务已经创建\n🌟任务 ID：{}\n✨prompt: {}".format(mjtask.from_nickname, mjtask.id, mjtask.origin_prompt)
                        botActions.to_user = mjtask.from_id
                        botActions.status = "waiting"
                        botActions.save()
                        # 将状态改为已开始
                        mjtask.status = "start"
                        mjtask.save()

# 处理 mj 已开始任务并生成机器人消息
def process_start_tasks():
    global MJ_RECEIVER_ON
    while MJ_RECEIVER_ON:
        # 获取已经发送的 mj 任务
        mjtasks = MJTasks.objects.filter(Q(status="sent") | Q(status="start"))[:20]
        for mjtask in mjtasks:
            print(mjtask)
            # todo: 任务超过 X 分钟，将任务状态改为 done 并处理下一个
            if False:
                # 将状态改为已开始
                mjtask.status = "done"
                mjtask.save()
                continue
            extra = json.loads(mjtask.bot_info.extra.replace("'", "\""))
            # 创建一个接收器
            mjreceiver = Receiver(
                authorization=extra.get("authorization", MJ_DEFAULT_CONFIG["authorization"]),
                channelid=extra.get("channelid", MJ_DEFAULT_CONFIG["channelid"]),
                proxy="127.0.0.1:7890"
            )
            try:
                time.sleep(2)
                mjreceiver.collecting_results()
            except:
                # 有异常则跳过
                continue
            else:
                # 已完成任务逻辑
                for df_id in mjreceiver.df.index:
                    image_url = mjreceiver.df.loc[df_id][2]
                    # todo: 将 image_url 传到对象存储中，生成 oss_image_url 用于返回给客户端
                    # 将文件下载到本地，并转码
                    pic_res = requests.get(
                        url=image_url,
                        proxies={'http': 'http://127.0.0.1:7890', 'https': 'http://127.0.0.1:7890'}
                    )
                    origin_file_path = os.path.join("/tmp", mjreceiver.df.loc[df_id][3])
                    with open(origin_file_path, "wb") as f:
                        f.write(pic_res.content)
                    png_file_path = origin_file_path
                    if image_url.endswith("webp"):
                        png_file_path = os.path.join("/tmp", mjreceiver.df.loc[df_id][3].replace("webp", "png"))
                        image = ImageOps.exif_transpose(Image.open(origin_file_path))
                        image.save(png_file_path)
                    # 将文件上传到存储桶中
                    oss_image_url = cos.get_url(
                        bucket=COS_CONFIG["bucket"],
                        key=cos.upload_file(bucket=COS_CONFIG["bucket"],key="wxmjbot/"+ mjreceiver.df.loc[df_id][3], local_path=png_file_path)
                    )

                    mj_done_msg = mjreceiver.df.loc[df_id][5]

                    if mjreceiver.df.loc[df_id][5].get('type') == 19:
                        # 处理变换任务
                        if "Variations by" in mj_done_msg['content']:
                            # 如果 origin_prompt 长度是 2，且包含 V，则生成回复
                            if len(mjtask.origin_prompt) == 2 and "V" in mjtask.origin_prompt:
                                # 生成机器人消息，并将任务状态置为已完成
                                # 生成文本消息
                                botActions = BotActions()
                                botActions.bot_info = mjtask.bot_info
                                botActions.action_name = "send_text"
                                botActions.action_content = "@{} \n✅变换任务已经完成\n🌟变换任务 ID：{}\n🪄 变换：\n[ U1 ] [ U2 ] [ U3 ] [ U4 ] \n[ V1 ] [ V2 ] [ V3 ] [ V4 ] \n✏️ 可使用 [/up-任务ID-操作] 进行变换（U：取图，V：变化）".format(mjtask.from_nickname, mjtask.id)
                                botActions.to_user = mjtask.from_id
                                botActions.status = "waiting"
                                botActions.save()
                                # 生成图片消息
                                botActions = BotActions()
                                botActions.bot_info = mjtask.bot_info
                                botActions.action_name = "send_image"
                                botActions.action_content = oss_image_url
                                botActions.to_user = mjtask.from_id
                                botActions.status = "waiting"
                                botActions.save()
                                # 将状态改为已完成
                                mjtask.status = "done"
                                mjtask.origin_image_url = image_url
                                mjtask.oss_image_url = oss_image_url
                                mjtask.mj_done_msg = json.dumps(mj_done_msg)
                                mjtask.save()
                        else:
                            # 根据 UX 返回对应数据
                            if mjtask.origin_prompt == "U{}".format(mj_done_msg['content'].split("Image #")[1][0]):
                                # 生成机器人消息，并将任务状态置为已完成
                                # 生成文本消息
                                botActions = BotActions()
                                botActions.bot_info = mjtask.bot_info
                                botActions.action_name = "send_text"
                                botActions.action_content = "@{} \n🎨 绘画成功!\n🌟变换任务 ID：{}".format(mjtask.from_nickname, mjtask.id)
                                botActions.to_user = mjtask.from_id
                                botActions.status = "waiting"
                                botActions.save()
                                # 生成图片消息
                                botActions = BotActions()
                                botActions.bot_info = mjtask.bot_info
                                botActions.action_name = "send_image"
                                botActions.action_content = oss_image_url
                                botActions.to_user = mjtask.from_id
                                botActions.status = "waiting"
                                botActions.save()
                                # 将状态改为已完成
                                mjtask.status = "done"
                                mjtask.origin_image_url = image_url
                                mjtask.oss_image_url = oss_image_url
                                mjtask.mj_done_msg = json.dumps(mj_done_msg)
                                mjtask.save()

                    else:
                        # 信息初始化，处理已完成任务
                        prompt = mjreceiver.df.loc[df_id][0].replace("'", " ").strip()[0:1500]
                        if mjtask.trans_prompt == prompt:
                            # 生成机器人消息，并将任务状态置为已完成
                            # 生成文本消息
                            botActions = BotActions()
                            botActions.bot_info = mjtask.bot_info
                            botActions.action_name = "send_text"
                            botActions.action_content = "@{} \n✅任务已经完成\n🌟任务 ID：{}\n✨prompt: {}\n🪄 变换：\n[ U1 ] [ U2 ] [ U3 ] [ U4 ] \n[ V1 ] [ V2 ] [ V3 ] [ V4 ] \n✏️ 可使用 [/up-任务ID-操作] 进行变换（U：取图，V：变化）\n🔗图片链接：{}".format(mjtask.from_nickname, mjtask.id, mjtask.origin_prompt, oss_image_url)
                            botActions.to_user = mjtask.from_id
                            botActions.status = "waiting"
                            botActions.save()
                            # 生成图片消息
                            botActions = BotActions()
                            botActions.bot_info = mjtask.bot_info
                            botActions.action_name = "send_image"
                            botActions.action_content = oss_image_url
                            botActions.to_user = mjtask.from_id
                            botActions.status = "waiting"
                            botActions.save()
                            # 将状态改为已完成
                            mjtask.status = "done"
                            mjtask.origin_image_url = image_url
                            mjtask.oss_image_url = oss_image_url
                            mjtask.mj_done_msg = json.dumps(mj_done_msg)
                            mjtask.save()

# 启动 mj 接收器相关任务
def start_mj_receiver(request):
    global MJ_RECEIVER_ON
    MJ_RECEIVER_ON = not MJ_RECEIVER_ON
    # 启动线程
    Thread(target=process_sent_tasks).start()
    Thread(target=process_start_tasks).start()
    return JsonResponse({
        "code": 0
    })

# 处理机器人消息并生成机器人行为
def generate_actions():
    global GENERATE_ACTIONS_ON
    while GENERATE_ACTIONS_ON:
        botMessages = BotMessage.objects.filter(state="waiting")[:10]
        for botMessage in botMessages:
            time.sleep(2)
            print("性能优化的地儿。。。。。。")
            if botMessage.bot_info.bot_type == "wx_gpt":
                extra = json.loads(botMessage.bot_info.extra.replace("'", "\""))
                # 如果是微信 gpt 机器人
                print("wx_gpt")
                # 群聊消息处理
                if botMessage.source == "group":
                    origin_message = json.loads(botMessage.message)
                    # 群聊白名单需要放到数据库的 bot 信息中
                    group_name = origin_message["User"]["NickName"]
                    group_name_white_list = extra.get("group_name_white_list", "")
                    if group_name_white_list != None and 'ALL_GROUP' in group_name_white_list:
                        pass
                    elif group_name_white_list != None and group_name in group_name_white_list:
                        pass
                    else:
                        # 将消息的状态置成已忽略
                        botMessage.state = "ignore"
                        botMessage.save()
                        continue

                    # 如果不是 at 的消息
                    if not origin_message["IsAt"]:
                        # 将消息的状态置成已忽略
                        botMessage.state = "ignore"
                        botMessage.save()
                        continue
                    else:
                        # at 消息处理逻辑
                        # 敏感词提示
                        for sensitiveWord in sensitiveWords:
                            print("开始敏感词检测")
                            if sensitiveWord.word in origin_message["Content"]:
                                # 生成机器人行为
                                botActions = BotActions()
                                botActions.bot_info = botMessage.bot_info
                                botActions.action_name = "send_text"
                                botActions.action_content = "❌(•'╻'• ۶)۶当前prompt检测到敏感词！！！"
                                botActions.to_user = origin_message['FromUserName']
                                botActions.status = "waiting"
                                botActions.save()
                                # 将消息的状态置成已忽略
                                botMessage.state = "ignore"
                                botMessage.save()
                                continue
                        # 无敏感词，则请求 GPT 并生成回复的内容
                        # 生成机器人行为
                        botActions = BotActions()
                        botActions.bot_info = botMessage.bot_info
                        botActions.action_name = "send_text"
                        botActions.action_content = "gpt 功能完善中"
                        botActions.to_user = origin_message['FromUserName']
                        botActions.status = "waiting"
                        botActions.save()
                        # 将消息的状态置成已忽略
                        botMessage.state = "done"
                        botMessage.save()
                        continue

                # 私聊消息处理
                if botMessage.source == "single":
                    pass

            if botMessage.bot_info.bot_type == "wx_mj":
                # 如果是微信 mj 机器人，发送响应并创建 MJTasks
                print("wx_mj")
                # 如果是群聊
                if botMessage.source == "group":
                    extra = json.loads(botMessage.bot_info.extra.replace("'", "\""))
                    origin_message = json.loads(botMessage.message)
                    # 群聊白名单需要放到数据库的 bot 信息中
                    group_name = origin_message["User"]["NickName"]
                    group_name_white_list = extra.get("group_name_white_list", "")
                    if group_name_white_list != None and 'ALL_GROUP' in group_name_white_list:
                        pass
                    elif group_name_white_list != None and group_name in group_name_white_list:
                        pass
                    else:
                        # 将消息的状态置成已忽略
                        botMessage.state = "ignore"
                        botMessage.save()
                        continue

                    if origin_message["Content"].startswith("/imagine"):
                        # 敏感词提示
                        for sensitiveWord in sensitiveWords:
                            print("开始敏感词检测")
                            if sensitiveWord.word in origin_message["Content"]:
                                # 生成机器人行为
                                botActions = BotActions()
                                botActions.bot_info = botMessage.bot_info
                                botActions.action_name = "send_text"
                                botActions.action_content = "❌(•'╻'• ۶)۶当前prompt检测到敏感词！！！"
                                botActions.to_user = origin_message['FromUserName']
                                botActions.status = "waiting"
                                botActions.save()
                                # 将消息的状态置成已忽略
                                botMessage.state = "ignore"
                                botMessage.save()
                                return None
                        # 无敏感词，执行 prompt 相关逻辑
                        origin_prompt = origin_message["Content"].replace("/imagine", "").replace("'", " ").replace("\n", " ")
                        pic_url_list = re.findall( r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', origin_prompt)
                        if len(pic_url_list) > 0:
                            # 如果是图生图，将图片链接过滤掉再翻译
                            origin_prompt_without_pic = origin_prompt
                            # 如果是图生图，将图片链接过滤掉再翻译
                            for pic_url in pic_url_list:
                                origin_prompt_without_pic = origin_prompt_without_pic.replace(pic_url, "")
                            # 获取 trans_prompt
                            trans_prompt = baidu.translate(origin_prompt_without_pic.split('--')[0]).get("result").get('trans_result')[0]['dst'].replace("'", " ").strip()[0:1500]
                            trans_prompt = re.sub(r'[^a-zA-Z0-9\s]+', '', trans_prompt).lower()
                            # 加参数 加图片链接
                            final_prompt = " ".join(pic_url_list) + " " + trans_prompt + " " + origin_prompt.replace(origin_prompt.split('--')[0], "")
                        else:
                            # 如果不是图生图
                            # 获取 trans_prompt
                            trans_prompt = baidu.translate(origin_prompt.split('--')[0]).get("result").get('trans_result')[0]['dst'].replace("'", " ").strip()[0:1500]
                            trans_prompt = re.sub(r'[^a-zA-Z0-9\s]+', '', trans_prompt).lower()
                            final_prompt = trans_prompt + " " + origin_prompt.replace(origin_prompt.split('--')[0], "")
                        # 定义 mjsender
                        mjsender = Sender(
                            authorization=extra.get("authorization", MJ_DEFAULT_CONFIG["authorization"]),
                            channelid=extra.get("channelid", MJ_DEFAULT_CONFIG["channelid"]),
                            application_id = extra.get("application_id", MJ_DEFAULT_CONFIG["application_id"]),
                            guild_id = extra.get("guild_id", MJ_DEFAULT_CONFIG["guild_id"]),
                            session_id = extra.get("session_id", MJ_DEFAULT_CONFIG["session_id"]),
                            version = extra.get("version", MJ_DEFAULT_CONFIG["version"]),
                            id = extra.get("id", MJ_DEFAULT_CONFIG["id"]),
                            flags = extra.get("flags", MJ_DEFAULT_CONFIG["flags"]),
                            proxy = "127.0.0.1:7890"
                        )
                        try:
                            mjsender.send(final_prompt)
                            # 生成 MJ 任务
                            mjtasks = MJTasks()
                            mjtasks.bot_info = botMessage.bot_info
                            mjtasks.origin_prompt = origin_prompt
                            mjtasks.trans_prompt = trans_prompt
                            mjtasks.from_id = origin_message['FromUserName']
                            mjtasks.from_nickname = origin_message['ActualNickName']
                            mjtasks.status = "sent"
                            mjtasks.save()
                            print(mjtasks.id)
                            # todo: 生成机器人行为
                            # botActions = BotActions()
                            # botActions.bot_info = botMessage.bot_info
                            # botActions.action_name = "send_text"
                            # botActions.action_content = "@{} \n🚀任务已经发送至 Midjourney\n🌟任务 ID：{}\n✨prompt: {}".format(origin_message['ActualNickName'], mjtasks.id, origin_prompt),
                            # botActions.to_user = origin_message['FromUserName']
                            # botActions.status = "waiting"
                            # botActions.save()
                            # 将消息的状态置成已完成
                            botMessage.state = "done"
                            botMessage.save()
                            continue
                        except Exception:
                            # 等待后续优化
                            print("有异常", Exception.with_traceback())
                            # 将消息的状态置成已完成
                            botMessage.state = "done"
                            botMessage.save()
                            continue

                # 如果是私聊
                if botMessage.source == "single":
                    pass

# 开始根据消息生成行为
def start_generate_actions(request):
    global GENERATE_ACTIONS_ON
    GENERATE_ACTIONS_ON = not GENERATE_ACTIONS_ON
    Thread(target=generate_actions).start()
    return JsonResponse({
        "code": 0
    })

# 机器人未完成的行为获取接口
def get_waiting_actions(request):
    # 获取密钥信息
    api_key = request.headers.get("X-Api-Key")
    # 查询密钥信息
    try:
        botApiKey = BotApiKey.objects.get(api_key=api_key)
        botInfo = BotInfo.objects.get(api_key=botApiKey)
    except BotApiKey.DoesNotExist:
        # 如果密钥不存在，则返回错误信息
        return JsonResponse({
            "code": 400,
            "message": "无相关密钥或是密钥"
        })
    except BotInfo.DoesNotExist:
        # 如果机器人信息不存在，则返回错误信息
        return JsonResponse({
            "code": 400,
            "message": "机器人未注册，请联系管理员"
        })
    else:
        botActions = BotActions.objects.filter(status="waiting")
        print(botActions)
        result = []
        for botAction in botActions:
            if botAction.bot_info.api_key.api_key == api_key:
                result.append({
                    "id": botAction.id,
                    "action_name": botAction.action_name,
                    "to_user": botAction.to_user,
                    "action_content": botAction.action_content.replace("\\n", "\n").replace("('", "").replace("',)", "")
                })
        return JsonResponse({
            "code": 0,
            "result": result
        })

# 机器人行为完成接口
def set_actions_done(request):
    # 获取密钥信息
    api_key = request.headers.get("X-Api-Key")
    # 查询密钥信息
    try:
        botApiKey = BotApiKey.objects.get(api_key=api_key)
        botInfo = BotInfo.objects.get(api_key=botApiKey)
    except BotApiKey.DoesNotExist:
        # 如果密钥不存在，则返回错误信息
        return JsonResponse({
            "code": 400,
            "message": "无相关密钥或是密钥已经被激活"
        })
    except BotInfo.DoesNotExist:
        # 如果机器人信息不存在，则返回错误信息
        return JsonResponse({
            "code": 400,
            "message": "机器人未注册，请联系管理员"
        })
    else:
        botAction = BotActions.objects.get(id=request.GET.get("id"))
        botAction.status = "done"
        botAction.save()
        return JsonResponse({
            "code": 0
        })

